from .jagged_ops import jagged_log_softmax